[None][feat] Enable PyTorch profiler traces for VisualGen#16814
Draft
chang-l wants to merge 4 commits into
Draft
[None][feat] Enable PyTorch profiler traces for VisualGen#16814chang-l wants to merge 4 commits into
chang-l wants to merge 4 commits into
Conversation
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Contributor
WalkthroughVisualGen profiling now supports environment-controlled PyTorch traces coordinated with CUDA profiler ranges, denoising-phase hooks, rank-specific trace output, executor integration through ChangesVisualGen profiling and inference integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant DiffusionExecutor
participant BasePipeline
participant QwenImagePipeline
participant CUDAProfiler
participant TorchProfiler
DiffusionExecutor->>BasePipeline: Call run_inference(req)
BasePipeline->>CUDAProfiler: Start configured profiling window
BasePipeline->>TorchProfiler: Start tracing
BasePipeline->>QwenImagePipeline: Execute generation phases and denoise steps
QwenImagePipeline->>BasePipeline: Signal denoise and step boundaries
BasePipeline->>CUDAProfiler: Stop profiling window
BasePipeline->>TorchProfiler: Stop and export rank-specific trace
BasePipeline-->>DiffusionExecutor: Return PipelineOutput
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/pipeline.py (1)
191-212: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate profiler lifecycle methods.
Add
-> Noneto both methods.- def _cuda_profiler_start(self): + def _cuda_profiler_start(self) -> None: ... - def _cuda_profiler_stop(self): + def _cuda_profiler_stop(self) -> None:As per coding guidelines, “Annotate every function, use
Nonefor non-returning functions.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/pipeline.py` around lines 191 - 212, Add the `-> None` return annotation to both `_cuda_profiler_start` and `_cuda_profiler_stop`, preserving their existing profiler lifecycle behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 178-179: Update the trace path construction around
self._torch_profile_trace_path so each profiling window receives a unique output
filename instead of reusing the same rank-based path. Add and maintain a window
index across repeated A-B,C-D profiling stops, incorporating it into the
exported trace path while preserving the existing rank and extension components.
In `@tests/unittest/_torch/visual_gen/test_profiler.py`:
- Around line 17-88: The profiler tests cover only a single start/stop window;
add a test for multiple configured ranges such as 0-1 and 3-4, asserting the
corresponding profiling behavior across both windows. Also register
tests/unittest/_torch/visual_gen/test_profiler.py in the relevant test-db and QA
test-list files, preserving their existing list format.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 191-212: Add the `-> None` return annotation to both
`_cuda_profiler_start` and `_cuda_profiler_stop`, preserving their existing
profiler lifecycle behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ce070dc7-3ebc-4306-babd-3cbb37cc95e9
📒 Files selected for processing (3)
docs/source/developer-guide/perf-analysis.mdtensorrt_llm/_torch/visual_gen/pipeline.pytests/unittest/_torch/visual_gen/test_profiler.py
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/visual_gen/pipeline.py (2)
181-190: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSelect the profiler activity at runtime.
torch.profiler.profile()should not be given bothCUDAandXPUunconditionally; on builds that only support one accelerator backend, profiler setup can fail. Build the list from the supported activities or branch on the active device so the trace requests only CPU + the matching accelerator.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/pipeline.py` around lines 181 - 190, Update the profiler setup around self._torch_profiler and activities to select only the accelerator activity supported by the active device or build, always retaining CPU while excluding the incompatible CUDA/XPU activity. Pass the resulting CPU-plus-matching-accelerator list to torch.profiler.profile.
189-189: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
with_modules=Truewon’t record module hierarchy here — VisualGen uses eagernn.Modules, so the profiler trace won’t include module info unless this path is scripted/traced or you add explicit instrumentation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/pipeline.py` at line 189, The profiler configuration using with_modules=True does not capture hierarchy for VisualGen’s eager nn.Module execution. Remove this ineffective option or replace it with explicit module instrumentation, ensuring module hierarchy is recorded only through a supported mechanism.
🧹 Nitpick comments (1)
tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py (1)
226-253: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExpand profiler lifecycle coverage
test_forward_honors_profile_step_rangeonly covers the0-1path. Add focused cases forpredenoise,all,postdenoise, warmup gating, and the final inactive state; thepostdenoisecase should assert that tracing stops and exports.Coverage summary: added
test_forward_honors_profile_step_range; listed intests/integration/test_lists/test-db/l0_a10.yml; no matching entry intests/integration/test_lists/qa/. Coverage verdict: needs follow-up.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py` around lines 226 - 253, Expand test_forward_honors_profile_step_range with focused profiler lifecycle cases covering predenoise, all, postdenoise, warmup gating, and the final inactive state. Configure each profile range and assert the corresponding CUDA and torch profiler calls, especially that postdenoise stops tracing and exports the configured trace. Keep the tests isolated with the existing pipeline test doubles and mocks.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py`:
- Around line 528-530: Update the forward path around
_start_postdenoise_profile() so post-denoise profiling is stopped and exported
after _decode_latents() completes. Use a finally block to call the corresponding
stop operation even when decoding raises, ensuring profiler state is not left
active and the Chrome trace is exported.
- Around line 488-494: Move the `_start_predenoise_profile()` call earlier in
the request flow, before timestep creation and scheduler preparation, so
pre-loop scheduler refresh work is included in profiling. Keep
`_start_denoise_profile()` immediately before the denoising loop.
---
Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 181-190: Update the profiler setup around self._torch_profiler and
activities to select only the accelerator activity supported by the active
device or build, always retaining CPU while excluding the incompatible CUDA/XPU
activity. Pass the resulting CPU-plus-matching-accelerator list to
torch.profiler.profile.
- Line 189: The profiler configuration using with_modules=True does not capture
hierarchy for VisualGen’s eager nn.Module execution. Remove this ineffective
option or replace it with explicit module instrumentation, ensuring module
hierarchy is recorded only through a supported mechanism.
---
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py`:
- Around line 226-253: Expand test_forward_honors_profile_step_range with
focused profiler lifecycle cases covering predenoise, all, postdenoise, warmup
gating, and the final inactive state. Configure each profile range and assert
the corresponding CUDA and torch profiler calls, especially that postdenoise
stops tracing and exports the configured trace. Keep the tests isolated with the
existing pipeline test doubles and mocks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dbb8a4b9-421d-4cb2-bc91-1fad5661b0e9
📒 Files selected for processing (4)
docs/source/developer-guide/perf-analysis.mdtensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.pytensorrt_llm/_torch/visual_gen/pipeline.pytests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/source/developer-guide/perf-analysis.md
Comment on lines
+528
to
+530
| self._stop_step_profile(i) | ||
|
|
||
| self._start_postdenoise_profile() |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Stop and export post-denoise profiling before returning.
When post-denoise profiling is pending, Line 530 starts the CUDA and PyTorch profilers, but this forward() path never stops them after _decode_latents(). Consequently, post-denoise mode does not export its Chrome trace and can leave profiler state active. Stop it after decoding, preferably in a finally block.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py`
around lines 528 - 530, Update the forward path around
_start_postdenoise_profile() so post-denoise profiling is stopped and exported
after _decode_latents() completes. Use a finally block to call the corresponding
stop operation even when decoding raises, ensuring profiler state is not left
active and the Chrome trace is exported.
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
chang-l
requested review from
QiJune,
brnguyen2,
crazydemo,
jieli-matrix,
mlefeb01 and
mzweilz
July 24, 2026 06:35
chang-l
marked this pull request as draft
July 24, 2026 16:38
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dev Engineer Review
TLLM_TORCH_PROFILE_TRACE, recording CPU/CUDA/XPU activity with shapes/module metadata and exporting rank-suffixed Chrome trace files (visual-gen-trace-rank-*.json, plus-window-Nwhen multiple capture windows are used).TLLM_PROFILE_VISUAL_GEN_START_STOP, including phase handling forpredenoise,postdenoise, andall, with “warmup” aware behavior.torch.profiler.profilesession and ensured trace export occurs only for active profiling windows._profile_denoise_start/_end,_start_step_profile/_stop_step_profile) around the denoise loop.pipeline.run_inference()(replacinginfer()in relevant call sites and failure-propagation tests).QA Engineer Review
Test functions added/updated
tests/unittest/_torch/visual_gen/test_profiler.pytest_setup_torch_profilertest_setup_torch_profiler_requires_profile_rangetest_cuda_profiler_controls_torch_profilertest_cuda_profiler_uses_fresh_trace_for_each_windowtest_cuda_profiler_closes_cuda_gate_when_trace_export_failstest_run_inference_owns_request_profile_boundaries(parameterized)test_run_inference_defers_predenoise_profile_until_after_warmuptests/unittest/_torch/visual_gen/test_qwen_image_pipeline.pytest_forward_honors_profile_step_rangetests/unittest/_torch/visual_gen/test_visual_gen_params.pytests/unittest/visual_gen/test_executor_shared_tensor_ipc.pyCoverage in test lists
tests/integration/test_lists/test-db/l0_a10.yml: addedunittest/_torch/visual_gen/test_profiler.pyunder thevisual_gensection.tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.pyor the other updated VisualGen unit tests.Verdict: insufficient —
test_profiler.pyis added totest-db, but the additional/updated VisualGen profiler-related tests don’t appear to be covered by the updated test-list entries.++ BrianLi23 for viz
Description
VisualGen already uses
TLLM_PROFILE_VISUAL_GEN_START_STOPto gate Nsight Systems capture, but it could not emit PyTorch/Kineto traces like the LLM executor.This change starts an optional
torch.profiler.profilesession on the same VisualGen ranges whenTLLM_TORCH_PROFILE_TRACEis set:allnow covers the complete user request from text encoding through denoising and VAE decode.predenoisecovers request start through the denoise-loop boundary;postdenoisecovers that boundary through request completion.BasePipeline.run_inference()wrapper owns request-level start/stop and exception cleanup. Base denoise and Qwen-Image's custom true-CFG loop call the same semantic phase/step hooks, so Qwen does not duplicate profiler setup or teardown.cudaProfilerStart()/cudaProfilerStop()remain capture gates for an attached Nsight collector; they do not start Nsight on their own. Torch and Nsight collectors should be run in separate invocations because both use CUPTI.Test Coverage
test_profiler.pyinl0_a10.yml.13 passed, 38 warnings in 143.25s.error=NoneusingTLLM_PROFILE_VISUAL_GEN_START_STOP=all.PR Checklist
PR description clearly explains what and why.
PR follows the TRT-LLM coding guidelines.
Test cases are provided for the new code paths and are routed in the test database.
No public API changes or new dependencies.
Documentation is updated.
No CODEOWNERS or architecture-diagram change is required.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, comment
/bot help.